#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
__author__ = 'Justin'
__mtime__ = '2020-12-17'

"""

import torch
import torch.nn as nn
from collections import OrderedDict
from torch.autograd import Variable

class Graphencoder(nn.Module):
    def __init__(self,):
        super(Graphencoder, self).__init__()

        # output_size = 1 + (input_size + 2*padding - kernel_size)/stride
        self.encoder = nn.Sequential(
            OrderedDict([
                ("deconv1", nn.ConvTranspose2d(28, 64, kernel_size=2, stride=2, padding=0)),  # 16n x 16n
                ("relu1", nn.ReLU()),
                ("conv1", nn.Conv2d(64, 128, kernel_size=2, stride=2, padding=0)),  # 8n x 8n
                ("relu2", nn.ReLU()),
                ("conv2", nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1)),  # 4n x 4n
                ("relu3", nn.ReLU()),
                ("conv3", nn.Conv2d(256, 128, kernel_size=3, stride=1, padding=1)),  # 2n x 2n
                ("relu3", nn.ReLU()),
                ("conv4", nn.Conv2d(128, 1, kernel_size=3, stride=1, padding=1)),  # n x n
                ("relu4", nn.ReLU()),
                # ("conv5", nn.Conv2d(64, 1, kernel_size=3, stride=1, padding=1)),  # n x n
                # ("relu5", nn.ReLU()),
            ])
        )

        self.sigmoid = nn.Sigmoid()

    # def __init__(self,):
    #     super(Graphencoder, self).__init__()
    #
    #     # output_size = 1 + (input_size + 2*padding - kernel_size)/stride
    #     self.encoder = nn.Sequential(
    #         OrderedDict([
    #             ("deconv1", nn.ConvTranspose2d(28, 32, kernel_size=2, stride=2, padding=0)),  # 2n x 2n
    #             ("drelu1", nn.ReLU()),
    #             ("deconv2", nn.ConvTranspose2d(32, 64, kernel_size=2, stride=2, padding=0)),  # 4n x 4n
    #             ("drelu2", nn.ReLU()),
    #             ("conv1", nn.Conv2d(64, 128, kernel_size=2, stride=2, padding=0)),  # 2n x 2n
    #             ("relu2", nn.ReLU()),
    #             ("conv2", nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1)),  # 2n x 2n
    #             ("relu3", nn.ReLU()),
    #             ("conv3", nn.Conv2d(256, 128, kernel_size=2, stride=2, padding=0)),  # n x n
    #             ("relu3", nn.ReLU()),
    #             ("conv4", nn.Conv2d(128, 1, kernel_size=3, stride=1, padding=1)),  # n x n
    #             ("relu4", nn.ReLU()),
    #             # ("conv5", nn.Conv2d(64, 1, kernel_size=3, stride=1, padding=1)),  # n x n
    #             # ("relu5", nn.ReLU()),
    #         ])
    #     )
    #
    #     self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        z = self.encoder(x)
        # out = self.sigmoid(z)
        return z

class Graphencoder2(nn.Module):
    def __init__(self,):
        super(Graphencoder2, self).__init__()

        # Conv2d output_size = 1 + (input_size + 2*padding - kernel_size)/stride
        # deconv2d output = (input -1)*stride + outputpadding - 2 * padding + kernelsize
        self.encoder = nn.Sequential(
            OrderedDict([
                ("deconv1", nn.ConvTranspose2d(28, 32, kernel_size=2, stride=2, padding=0)),  # 2n x 2n
                ("relu1", nn.ReLU()),
                ("conv1", nn.Conv2d(32, 16, kernel_size=3, stride=1, padding=1)),  # 2n x 2n
                ("relu2", nn.ReLU()),
                ("deconv2", nn.ConvTranspose2d(16, 16, kernel_size=2, stride=2, padding=0)),  # 4n x 4n
                ("relu3", nn.ReLU()),
                ("conv2", nn.Conv2d(16, 32, kernel_size=3, stride=1, padding=1)),  # 4n x 4n
                ("relu4", nn.ReLU()),
                ("conv3", nn.Conv2d(32, 16, kernel_size=2, stride=2, padding=0)),  # 2n x 2n
                ("relu5", nn.ReLU()),
                ("conv4", nn.Conv2d(16, 8, kernel_size=3, stride=1, padding=1)),  # 2n x 2n
                ("relu6", nn.ReLU()),
                ("conv5", nn.Conv2d(8, 4, kernel_size=2, stride=2, padding=0)),  # n x n
                ("relu7", nn.ReLU()),
                ("conv6", nn.Conv2d(4, 1, kernel_size=3, stride=1, padding=1)),  # n x n
                ("relu8", nn.ReLU()),
            ])
        )

        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        z = self.encoder(x)
        return z